package com.huangjue.backend.filter;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.huangjue.backend.utils.EncryptionUtil;
import org.springframework.stereotype.Component;
import org.springframework.util.StreamUtils;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpServletResponseWrapper;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

/**
 * 加密解密过滤器
 * 自动处理请求解密和响应加密
 * 
 * 功能：
 * 1. 检查请求是否加密
 * 2. 解密请求数据
 * 3. 加密响应数据
 * 4. 防止重放攻击
 */
@Component
public class EncryptionFilter implements Filter {
    
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    // 最大时间戳有效期（5分钟）
    private static final long MAX_TIMESTAMP_AGE = 5 * 60 * 1000;
    
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {
        
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;
        
        // 处理CORS预检请求
        if ("OPTIONS".equalsIgnoreCase(httpRequest.getMethod())) {
            httpResponse.setHeader("Access-Control-Allow-Origin", "*");
            httpResponse.setHeader("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS");
            httpResponse.setHeader("Access-Control-Allow-Headers", "*");
            httpResponse.setHeader("Access-Control-Max-Age", "3600");
            httpResponse.setStatus(HttpServletResponse.SC_OK);
            return;
        }
        
        // 为所有请求添加CORS头
        httpResponse.setHeader("Access-Control-Allow-Origin", "*");
        httpResponse.setHeader("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS");
        httpResponse.setHeader("Access-Control-Allow-Headers", "*");
        
        // 检查请求是否加密
        String encryptedHeader = httpRequest.getHeader("X-Encrypted");
        String timestampHeader = httpRequest.getHeader("X-Timestamp");
        
        if ("true".equals(encryptedHeader)) {
            // 处理加密请求
            handleEncryptedRequest(httpRequest, httpResponse, chain, timestampHeader);
        } else {
            // 处理普通请求
            chain.doFilter(request, response);
        }
    }
    
    /**
     * 处理加密请求
     */
    private void handleEncryptedRequest(HttpServletRequest request, HttpServletResponse response, 
                                      FilterChain chain, String timestampHeader) 
            throws IOException, ServletException {
        
        try {
            // 1. 读取请求体
            String requestBody = StreamUtils.copyToString(request.getInputStream(), StandardCharsets.UTF_8);
            
            // 2. 解析加密数据
            Map<String, Object> encryptedData = objectMapper.readValue(requestBody, Map.class);
            String encryptedDataStr = (String) encryptedData.get("encryptedData");
            String encryptedAESKey = (String) encryptedData.get("encryptedAESKey");
            Long timestamp = (Long) encryptedData.get("timestamp");
            
            // 3. 验证时间戳（防止重放攻击）
            if (timestamp != null && !EncryptionUtil.isValidTimestamp(timestamp, MAX_TIMESTAMP_AGE)) {
                response.setStatus(HttpServletResponse.SC_UNAUTHORIZED);
                response.getWriter().write("{\"success\":false,\"message\":\"请求已过期\"}");
                return;
            }
            
            // 4. 解密请求数据
            String decryptedData = EncryptionUtil.decryptRequestData(encryptedDataStr, encryptedAESKey);
            
            // 5. 创建新的请求包装器
            RequestWrapper requestWrapper = new RequestWrapper(request, decryptedData);
            
            // 6. 创建响应包装器
            ResponseWrapper responseWrapper = new ResponseWrapper(response);
            
            // 7. 继续处理请求
            chain.doFilter(requestWrapper, responseWrapper);
            
            // 8. 直接返回响应数据（不加密响应）
            String responseBody = responseWrapper.getContentAsString();
            
            // 9. 设置响应头
            response.setContentType("application/json;charset=UTF-8");
            response.setCharacterEncoding("UTF-8");
            
            // 10. 写入响应
            response.getWriter().write(responseBody);
            
        } catch (Exception e) {
            System.err.println("加密请求处理失败: " + e.getMessage());
            response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
            response.getWriter().write("{\"success\":false,\"message\":\"请求解密失败\"}");
        }
    }
    
    /**
     * 请求包装器
     * 用于替换请求体中的加密数据为解密后的数据
     */
    private static class RequestWrapper extends HttpServletRequestWrapper {
        private final String decryptedData;
        
        public RequestWrapper(HttpServletRequest request, String decryptedData) {
            super(request);
            this.decryptedData = decryptedData;
        }
        
        @Override
        public ServletInputStream getInputStream() throws IOException {
            return new ServletInputStream() {
                private final ByteArrayInputStream inputStream = 
                    new ByteArrayInputStream(decryptedData.getBytes(StandardCharsets.UTF_8));
                
                @Override
                public boolean isFinished() {
                    return inputStream.available() == 0;
                }
                
                @Override
                public boolean isReady() {
                    return true;
                }
                
                @Override
                public void setReadListener(ReadListener readListener) {
                    // 不需要实现
                }
                
                @Override
                public int read() throws IOException {
                    return inputStream.read();
                }
            };
        }
        
        @Override
        public BufferedReader getReader() throws IOException {
            return new BufferedReader(new InputStreamReader(getInputStream(), StandardCharsets.UTF_8));
        }
    }
    
    /**
     * 响应包装器
     * 用于捕获响应数据
     */
    private static class ResponseWrapper extends HttpServletResponseWrapper {
        private final ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        private final PrintWriter writer = new PrintWriter(new OutputStreamWriter(outputStream, StandardCharsets.UTF_8));
        
        public ResponseWrapper(HttpServletResponse response) {
            super(response);
            response.setCharacterEncoding("UTF-8");
        }
        
        @Override
        public ServletOutputStream getOutputStream() throws IOException {
            return new ServletOutputStream() {
                @Override
                public boolean isReady() {
                    return true;
                }
                
                @Override
                public void setWriteListener(WriteListener writeListener) {
                    // 不需要实现
                }
                
                @Override
                public void write(int b) throws IOException {
                    outputStream.write(b);
                }
            };
        }
        
        @Override
        public PrintWriter getWriter() throws IOException {
            return writer;
        }
        
        public String getContentAsString() throws IOException {
            writer.flush();
            return outputStream.toString(StandardCharsets.UTF_8.name());
        }
        
        @Override
        public void setCharacterEncoding(String charset) {
            super.setCharacterEncoding(charset);
        }
    }
} 